Flattening Every Memory Peak in Long-Context Mixture-of-Experts Training
抹平长上下文 MoE 训练的所有显存峰值
🧭 长上下文 MoE 训练显存优化
Salesforce AI Research
针对长上下文 MoE 训练四大显存瓶颈提出四种精确调度，将 120B-667B 模型训练扩展到 1M 上下文，吞吐最高提升 10.4 倍
📌 工作参考价值（4 点）
- 若在业务中训练长上下文 MoE（如融合用户长期行为序列、Agent 记忆的推荐/广告模型），优先采用固定 GPU working set 的调度思路，而不是只看平均显存；PipelinedLLEP 与 Ring-DTP 都能在不改精度下显著压峰值。
- OffloadStreamAdamW 将 CPU optimizer offload 的串行 Adam 改为 bucket 流水线，能直接加速大规模 embedding/排序模型训练中的 offload 更新；若你已用 CPU offload 放优化器状态，可参照其分桶流水实现。
- SCO 只把 checkpoint 边界最占显存的长生命周期张量 offload 到 CPU，适合处理用户长序列的生成式推荐/Agent 轨迹训练，能在基本不增加复杂度的前提下扩大 context 长度。
- 这些方法只改变执行顺序和搬运粒度，梯度与 loss 严格一致，便于在 FSDP/Megatron 等框架上增量实现，风险低。